import serial
import time
import argparse
from PIL import Image

# 读取图像并将像素数据通过UART发送到BRAM

def main():
    parser = argparse.ArgumentParser(description="通过UART发送MNIST图像到FPGA BRAM")
    parser.add_argument('--port', type=str, required=True, help="串口号，例如 COM3")
    parser.add_argument('--baudrate', type=int, default=115200, help="波特率，默认为115200")
    parser.add_argument('--image', type=str, required=True, help="要发送的PNG图像路径")
    args = parser.parse_args()

    # 打开并转换为灰度图像
    img = Image.open(args.image).convert('L')
    pixels = list(img.getdata())  # 获取所有像素值
    if len(pixels) != 28*28:
        print(f"警告：图像尺寸为{img.size}，不是28x28")

    # 打开串口并发送数据
    with serial.Serial(args.port, args.baudrate, timeout=1) as ser:
        time.sleep(2)  # 等待FPGA上电和UART初始化
        # 发送帧头
        ser.write(b'START')
        # 依次发送像素值（0-255）
        for p in pixels:
            ser.write(bytes([p]))
        # 发送帧尾
        ser.write(b'END')
        ser.flush()

        # 清空串口缓存
        ser.reset_input_buffer()
        ser.reset_output_buffer()
        # 给 FPGA 一些处理时间
        time.sleep(0.5)
        # 循环读取 ACK
        ack = b''
        timeout = 5  # 秒
        start = time.time()
        while time.time() - start < timeout:
            c = ser.read(1)
            if c:
                ack += c
                if ack.endswith(b'DONE'):
                    print("FPGA 已成功接收并写入 BRAM")
                    break
        else:
            print(f"未收到ACK，返回数据: {ack}")

    print(f"已通过UART发送 {len(pixels)} 字节图像数据")

if __name__ == '__main__':
    main()
